import cv2, time, datetime
from pathlib import Path
from .logger import get_logger

class Camera:
    def __init__(self, cfg):
        self.log = get_logger("Camera")
        self.url = cfg["rtsp"]
        self.save_video = cfg["save_video"]
        self.save_img = cfg["save_img"]
        self.out_dir = Path(cfg["out_dir"])
        self.out_dir.mkdir(exist_ok=True)

        # 连接并保存首帧用于宽高、fps
        self._connect()

        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS)) or 25
        self.w = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.h = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.frame_id = 0

    # ------------- 内部：建立/重连 -------------
    def _connect(self):
        self.log.info("正在连接 RTSP：%s", self.url)
        # 强制 TCP + 超时
        tcp_url = f"{self.url}?transport=tcp"
        self.cap = cv2.VideoCapture(tcp_url, cv2.CAP_FFMPEG)
        self.cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 3000)
        self.cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, 1000)
        if not self.cap.isOpened():
            raise RuntimeError("RTSP 连接失败")
        self.log.info("RTSP 连接成功")

    # ------------- 对外：不断产出帧 -------------
    def frames(self):
        while True:
            ret, frame = self.cap.read()
            if not ret:
                self.log.warning("取帧失败，开始重连…")
                self.cap.release()
                time.sleep(1)
                self._connect()
                continue
            self.frame_id += 1
            yield frame

    # ------------- 其余接口保持原样 -------------
    def save_frame(self, frame, prefix=""):
        if self.save_img:
            t = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            cv2.imwrite(str(self.out_dir / f"{prefix}_{t}.jpg"), frame)
            self.log.debug("已保存图片：%s", t)

    def start_record(self):
        if not self.save_video:
            return None
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        t = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.writer = cv2.VideoWriter(
            str(self.out_dir / f"{t}.mp4"),
            fourcc, self.fps, (self.w, self.h))
        return self.writer

    def stop_record(self):
        if hasattr(self, "writer") and self.writer is not None:
            self.writer.release()
            self.writer = None

    def __del__(self):
        self.stop_record()